import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 参数
input_size=1
output_size=1
num_epochs=60
learning_rate=0.01

# 测试数据
x_train = np.array([[3.3], [4.4], [5.5], [6.71], [6.93], [4.168], 
                    [9.779], [6.182], [7.59], [2.167], [7.042], 
                    [10.791], [5.313], [7.997], [3.1]], dtype=np.float32)

y_train = np.array([[1.7], [2.76], [2.09], [3.19], [1.694], [1.573], 
                    [3.366], [2.596], [2.53], [1.221], [2.827], 
                    [3.465], [1.65], [2.904], [1.3]], dtype=np.float32)

# 线性模型
model=nn.Linear(input_size,output_size)

# 误差与优化
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    # 转换numpy数组成tensor
    inputs=torch.from_numpy(x_train)
    targets=torch.from_numpy(y_train)

    # 前向计算
    outputs=model(inputs)
    loss=criterion(outputs,targets)

    # 返回值计算和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch+1) % 5 ==0:
        print('训练集:[{}/{}],Loss:{:.4f}'.format(epoch+1,num_epochs,loss.item()))

# 出图
predicted=model(torch.from_numpy(x_train)).detach().numpy()
plt.plot(x_train,y_train,'ro',label='origin data')
plt.plot(x_train,predicted,label='fitted line')
plt.legend()
plt.show()

# 保存模型
#torch.save(model.state_dict(),'model.cpkt')

'''
训练集:[5/60],Loss:0.2166
训练集:[10/60],Loss:0.2154
训练集:[15/60],Loss:0.2142
训练集:[20/60],Loss:0.2131
训练集:[25/60],Loss:0.2119
训练集:[30/60],Loss:0.2108
训练集:[35/60],Loss:0.2098
训练集:[40/60],Loss:0.2087
训练集:[45/60],Loss:0.2077
训练集:[50/60],Loss:0.2067
训练集:[55/60],Loss:0.2058
训练集:[60/60],Loss:0.2048
'''